import numpy as np
import matplotlib.pyplot as plt
import sys
import os

def load_gbn_points(file_path):
    """加载 GBN 输出的点集文件"""
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"Point file not found: {file_path}")
    return np.loadtxt(file_path)

def compute_radial_spectrum(points, num_bins=100):
    """计算径向功率谱（论文图1(h)风格）"""
    num_points, dim = points.shape
    if dim != 2:
        raise ValueError("Spectrum calculation only supports 2D points")
    
    # 计算所有点对的距离
    dists = []
    for i in range(num_points):
        for j in range(i + 1, num_points):
            dx = points[i, 0] - points[j, 0]
            dy = points[i, 1] - points[j, 1]
            # 环形域距离（匹配论文中的边界处理）
            dx = min(dx, 1 - dx) if dx > 0.5 else max(dx, dx - 1)
            dy = min(dy, 1 - dy) if dy > 0.5 else max(dy, dy - 1)
            dist = np.sqrt(dx**2 + dy**2)
            dists.append(dist)
    
    # 统计距离分布（功率谱近似）
    dists = np.array(dists)
    bins = np.linspace(0, np.max(dists), num_bins)
    counts, _ = np.histogram(dists, bins=bins, density=True)
    return bins[:-1], counts

def plot_gbn_results(points, spectrum_bins, spectrum_counts, output_dir):
    """生成论文风格的可视化图"""
    plt.rcParams.update({
        "font.family": "serif",
        "font.size": 10,
        "axes.linewidth": 0.8,
        "figure.dpi": 300,
        "savefig.dpi": 300
    })

    # 创建 1x2 子图（点集 + 功率谱）
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

    # 子图1：2D 蓝噪声点集（论文图1(a)-(g)风格）
    ax1.scatter(points[:, 0], points[:, 1], s=0.8, c='black', alpha=0.8)
    ax1.set_xlim(0, 1)
    ax1.set_ylim(0, 1)
    ax1.set_aspect('equal')
    ax1.set_title(f"GBN Point Distribution (N={points.shape[0]})", pad=10)
    ax1.set_xlabel("X Coordinate")
    ax1.set_ylabel("Y Coordinate")
    ax1.grid(False)

    # 子图2：径向功率谱（论文图1(h)风格，对数坐标）
    ax2.semilogy(spectrum_bins, spectrum_counts, color='darkblue', linewidth=1.2)
    ax2.set_xlabel("Radial Frequency (Normalized Distance)")
    ax2.set_ylabel("Power Spectrum (Log Scale)")
    ax2.set_title("Radial Power Spectrum of GBN", pad=10)
    ax2.grid(True, alpha=0.3, linestyle='--')
    ax2.set_xlim(0, np.max(spectrum_bins))

    # 调整布局并保存
    plt.tight_layout()
    output_path = os.path.join(output_dir, "gbn_visualization.png")
    plt.savefig(output_path, bbox_inches='tight', dpi=300)
    print(f"Visualization saved to: {output_path}")

if __name__ == "__main__":
    # 命令行参数：点集文件路径 （输出目录为当前目录）
    #if len(sys.argv) != 2:
        #print(f"Usage: {sys.argv[0]} <gbn_points_file> ")
        #sys.exit(1)

    #points_file = sys.argv[1]
    #output_dir = "./output"
    #os.makedirs(output_dir, exist_ok=True)

     # 固定点集文件路径为上上一级目录的 build/bin/gbn_points_2d.txt
    points_file = os.path.abspath(os.path.join("..", "..","build", "bin", "gbn_points_2d.txt"))
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)

    # 执行可视化流程
    try:
        points = load_gbn_points(points_file)
        print(f"Loaded {points.shape[0]} points (dim={points.shape[1]})")
        
        if points.shape[1] == 2:
            spectrum_bins, spectrum_counts = compute_radial_spectrum(points)
            plot_gbn_results(points, spectrum_bins, spectrum_counts, output_dir)
        else:
            print(f"Warning: {points.shape[1]}D points not supported for visualization (only 2D)")
    except Exception as e:
        print(f"Visualization failed: {str(e)}")
        sys.exit(1)